import pygame
import sys
import numpy as np
import random
from pygame.locals import *

# 初始化pygame
pygame.init()

# 游戏常量
BOARD_SIZE = 15  # 15x15的棋盘
GRID_SIZE = 40   # 每个格子的像素大小
MARGIN = 30       # 棋盘边距
PIECE_RADIUS = 18 # 棋子半径
WINDOW_SIZE = BOARD_SIZE * GRID_SIZE + 2 * MARGIN  # 窗口大小

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
BOARD_COLOR = (220, 179, 92)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
GRAY = (200, 200, 200)

# 创建游戏窗口
screen = pygame.display.set_mode((WINDOW_SIZE, WINDOW_SIZE + 60))
pygame.display.set_caption('五子棋 - 人类(黑) vs 电脑(白)')

# 游戏状态
board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=int)  # 0:空, 1:黑, 2:白
current_player = 1  # 1:人类(黑), 2:电脑(白)
game_over = False
winner = 0
move_history = []  # 记录每一步的落子位置
scores = []  # 记录每局游戏的得分

# 字体
font = pygame.font.SysFont('Arial', 24)
small_font = pygame.font.SysFont('Arial', 16)

def draw_board():
    """绘制棋盘"""
    screen.fill(BOARD_COLOR)
    
    # 绘制网格线
    for i in range(BOARD_SIZE):
        # 横线
        pygame.draw.line(screen, BLACK, 
                         (MARGIN, MARGIN + i * GRID_SIZE), 
                         (WINDOW_SIZE - MARGIN, MARGIN + i * GRID_SIZE), 2)
        # 竖线
        pygame.draw.line(screen, BLACK, 
                         (MARGIN + i * GRID_SIZE, MARGIN), 
                         (MARGIN + i * GRID_SIZE, WINDOW_SIZE - MARGIN), 2)
    
    # 绘制五个小黑点(天元、星位)
    dots = [(3, 3), (3, 11), (7, 7), (11, 3), (11, 11)]
    for dot in dots:
        pygame.draw.circle(screen, BLACK, 
                           (MARGIN + dot[0] * GRID_SIZE, MARGIN + dot[1] * GRID_SIZE), 5)
    
    # 绘制棋子
    for y in range(BOARD_SIZE):
        for x in range(BOARD_SIZE):
            if board[y][x] == 1:  # 黑子
                pygame.draw.circle(screen, BLACK, 
                                  (MARGIN + x * GRID_SIZE, MARGIN + y * GRID_SIZE), PIECE_RADIUS)
            elif board[y][x] == 2:  # 白子
                pygame.draw.circle(screen, WHITE, 
                                  (MARGIN + x * GRID_SIZE, MARGIN + y * GRID_SIZE), PIECE_RADIUS)
                pygame.draw.circle(screen, BLACK, 
                                  (MARGIN + x * GRID_SIZE, MARGIN + y * GRID_SIZE), PIECE_RADIUS, 1)
    
    # 绘制游戏状态信息
    status_area = pygame.Rect(0, WINDOW_SIZE, WINDOW_SIZE, 60)
    pygame.draw.rect(screen, GRAY, status_area)
    
    if game_over:
        if winner == 1:
            status_text = font.render("游戏结束! 人类(黑)获胜!", True, GREEN)
        elif winner == 2:
            status_text = font.render("游戏结束! 电脑(白)获胜!", True, RED)
        else:
            status_text = font.render("游戏结束! 平局!", True, BLUE)
    else:
        if current_player == 1:
            status_text = font.render("轮到: 人类(黑) 落子", True, BLACK)
        else:
            status_text = font.render("电脑(白) 思考中...", True, BLACK)
    
    screen.blit(status_text, (20, WINDOW_SIZE + 10))
    
    # 显示历史得分
    if scores:
        score_text = small_font.render(f"历史得分: 人类 {sum(s[0] for s in scores)} - 电脑 {sum(s[1] for s in scores)}", True, BLACK)
        screen.blit(score_text, (20, WINDOW_SIZE + 40))

def check_win(x, y, player):
    """检查是否有玩家获胜"""
    directions = [(1, 0), (0, 1), (1, 1), (1, -1)]  # 横、竖、斜、反斜
    
    for dx, dy in directions:
        count = 1  # 当前位置已经有一个棋子
        
        # 正向检查
        nx, ny = x + dx, y + dy
        while 0 <= nx < BOARD_SIZE and 0 <= ny < BOARD_SIZE and board[ny][nx] == player:
            count += 1
            nx += dx
            ny += dy
        
        # 反向检查
        nx, ny = x - dx, y - dy
        while 0 <= nx < BOARD_SIZE and 0 <= ny < BOARD_SIZE and board[ny][nx] == player:
            count += 1
            nx -= dx
            ny -= dy
        
        if count >= 5:
            return True
    
    return False

def evaluate_board():
    """评估当前棋盘状态，返回人类和电脑的得分"""
    human_score = 0
    computer_score = 0
    
    # 评分标准
    score_table = {
        1: 1,      # 活1
        2: 10,     # 活2
        3: 100,    # 活3
        4: 1000,   # 活4
        5: 10000,  # 五连
        'blocked1': 0.5,   # 死1
        'blocked2': 5,     # 死2
        'blocked3': 50,    # 死3
        'blocked4': 500    # 死4
    }
    
    # 检查所有行、列、对角线
    for y in range(BOARD_SIZE):
        for x in range(BOARD_SIZE):
            if board[y][x] == 0:  # 空位置
                continue
            
            # 四个方向: 横、竖、斜、反斜
            for dx, dy in [(1, 0), (0, 1), (1, 1), (1, -1)]:
                # 检查这个方向是否已经有评估过
                prev_x, prev_y = x - dx, y - dy
                if (0 <= prev_x < BOARD_SIZE and 0 <= prev_y < BOARD_SIZE and 
                    board[prev_y][prev_x] == board[y][x]):
                    continue  # 已经评估过这个方向
                
                count = 1
                blocked = False
                
                # 正向检查
                nx, ny = x + dx, y + dy
                while True:
                    if nx >= BOARD_SIZE or ny >= BOARD_SIZE or ny < 0:
                        blocked = True
                        break
                    if board[ny][nx] != board[y][x]:
                        if board[ny][nx] != 0:
                            blocked = True
                        break
                    count += 1
                    nx += dx
                    ny += dy
                
                # 反向检查
                nx, ny = x - dx, y - dy
                while True:
                    if nx < 0 or ny < 0 or ny >= BOARD_SIZE:
                        blocked = True
                        break
                    if board[ny][nx] != board[y][x]:
                        if board[ny][nx] != 0:
                            blocked = True
                        break
                    count += 1
                    nx -= dx
                    ny -= dy
                
                # 计算得分
                if board[y][x] == 1:  # 人类
                    if count >= 5:
                        human_score += score_table[5]
                    elif not blocked:
                        if count in score_table:
                            human_score += score_table[count]
                    else:
                        if count == 1:
                            human_score += score_table['blocked1']
                        elif count == 2:
                            human_score += score_table['blocked2']
                        elif count == 3:
                            human_score += score_table['blocked3']
                        elif count == 4:
                            human_score += score_table['blocked4']
                else:  # 电脑
                    if count >= 5:
                        computer_score += score_table[5]
                    elif not blocked:
                        if count in score_table:
                            computer_score += score_table[count]
                    else:
                        if count == 1:
                            computer_score += score_table['blocked1']
                        elif count == 2:
                            computer_score += score_table['blocked2']
                        elif count == 3:
                            computer_score += score_table['blocked3']
                        elif count == 4:
                            computer_score += score_table['blocked4']
    
    return human_score, computer_score

def computer_move():
    """电脑AI走棋"""
    global current_player, game_over, winner
    
    if game_over:
        return
    
    # 1. 首先检查电脑是否有四连可以完成(但根据规则不能赢)
    # 2. 检查人类是否有四连需要阻止
    # 3. 否则选择最佳防守位置
    
    best_score = -float('inf')
    best_moves = []
    
    for y in range(BOARD_SIZE):
        for x in range(BOARD_SIZE):
            if board[y][x] != 0:
                continue
            
            # 尝试在这个位置落子
            board[y][x] = 2
            
            # 检查电脑是否会赢(不允许)
            if check_win(x, y, 2):
                board[y][x] = 0  # 撤销落子
                continue  # 跳过这个位置，不能赢
            
            # 检查人类是否会赢(必须阻止)
            human_will_win = False
            for yy in range(BOARD_SIZE):
                for xx in range(BOARD_SIZE):
                    if board[yy][xx] != 0:
                        continue
                    board[yy][xx] = 1
                    if check_win(xx, yy, 1):
                        human_will_win = True
                        board[yy][xx] = 0
                        break
                    board[yy][xx] = 0
                if human_will_win:
                    break
            
            if human_will_win:
                # 必须阻止人类赢
                board[y][x] = 0  # 撤销落子
                board[yy][xx] = 2  # 在人类要赢的位置落子
                move_history.append((xx, yy))
                
                # 检查游戏是否结束
                if check_win(xx, yy, 2):
                    game_over = True
                    winner = 2
                return
            
            # 评估这个位置的得分
            human_score, computer_score = evaluate_board()
            score = computer_score - human_score * 1.2  # 更注重防守
            
            # 如果是中心位置，加分
            center_dist = abs(x - BOARD_SIZE//2) + abs(y - BOARD_SIZE//2)
            score += (BOARD_SIZE - center_dist) * 2
            
            # 如果是靠近已有棋子的位置，加分
            has_neighbor = False
            for dy in [-1, 0, 1]:
                for dx in [-1, 0, 1]:
                    if dx == 0 and dy == 0:
                        continue
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < BOARD_SIZE and 0 <= ny < BOARD_SIZE and board[ny][nx] != 0:
                        has_neighbor = True
                        break
                if has_neighbor:
                    break
            if has_neighbor:
                score += 10
            
            # 记录最佳走法
            if score > best_score:
                best_score = score
                best_moves = [(x, y)]
            elif score == best_score:
                best_moves.append((x, y))
            
            board[y][x] = 0  # 撤销落子
    
    # 如果没有合法的走法(理论上不会发生)
    if not best_moves:
        game_over = True
        return
    
    # 随机选择一个最佳走法
    x, y = random.choice(best_moves)
    board[y][x] = 2
    move_history.append((x, y))
    
    # 检查游戏是否结束
    if check_win(x, y, 2):
        game_over = True
        winner = 2
    
    current_player = 1  # 轮到人类

def reset_game():
    """重置游戏"""
    global board, current_player, game_over, winner, move_history
    board = np.zeros((BOARD_SIZE, BOARD_SIZE), dtype=int)
    current_player = 1
    game_over = False
    winner = 0
    move_history = []
    
    # 计算并记录上一局的得分
    if len(move_history) > 0:
        human_score, computer_score = evaluate_board()
        scores.append((human_score, computer_score))

def main():
    """主游戏循环"""
    global current_player, game_over, winner
    
    clock = pygame.time.Clock()
    
    while True:
        for event in pygame.event.get():
            if event.type == QUIT:
                pygame.quit()
                sys.exit()
            
            elif event.type == MOUSEBUTTONDOWN and not game_over and current_player == 1:
                # 人类玩家落子
                x, y = event.pos[0], event.pos[1]
                
                # 转换为棋盘坐标
                grid_x = round((x - MARGIN) / GRID_SIZE)
                grid_y = round((y - MARGIN) / GRID_SIZE)
                
                # 检查是否在棋盘内且位置为空
                if (0 <= grid_x < BOARD_SIZE and 0 <= grid_y < BOARD_SIZE and 
                    board[grid_y][grid_x] == 0):
                    board[grid_y][grid_x] = 1
                    move_history.append((grid_x, grid_y))
                    
                    # 检查游戏是否结束
                    if check_win(grid_x, grid_y, 1):
                        game_over = True
                        winner = 1
                    else:
                        current_player = 2  # 轮到电脑
                        
                        # 延迟一下，让玩家看到自己的落子
                        pygame.time.delay(300)
                        computer_move()
            
            elif event.type == KEYDOWN:
                if event.key == K_r:  # 按R键重置游戏
                    reset_game()
                elif event.key == K_ESCAPE:  # 按ESC退出
                    pygame.quit()
                    sys.exit()
        
        # 绘制游戏
        draw_board()
        pygame.display.update()
        clock.tick(30)

if __name__ == "__main__":
    main()
